// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/container/f14_set.h>
#include "melon/hash/hash.h"

#include <pollux/expression/vector_function.h>
#include <pollux/functions/macros.h>
#include <pollux/functions/lib/registration_helpers.h>
#include <pollux/functions/sparksql/arena.h>
#include <pollux/functions/sparksql/comparisons.h>
#include <pollux/type/filter.h>

namespace kumo::pollux::functions::sparksql {
    namespace {
        template<typename T>
        class Set : public melon::F14FastSet<T, melon::hasher<T>, Equal<T> > {
        };

        template<>
        class Set<float> : public melon::F14FastSet<
                    float,
                    util::floating_point::NaNAwareHash<float>,
                    util::floating_point::NaNAwareEquals<float> > {
        };

        template<>
        class Set<double> : public melon::F14FastSet<
                    double,
                    util::floating_point::NaNAwareHash<double>,
                    util::floating_point::NaNAwareEquals<double> > {
        };

        template<>
        class Set<StringView> {
        public:
            using value_type = std::string_view;

            void emplace(const StringView &s) {
                std::string_view sv(s.data(), s.size());
                if (!set_.contains(sv)) {
                    set_.emplace(arena_.writeString(sv));
                }
            }

            bool contains(const StringView &s) const {
                return set_.contains(std::string_view(s.data(), s.size()));
            }

            void reserve(size_t size) {
                set_.reserve(size);
            }

            size_t size() const {
                return set_.size();
            }

            auto begin() const {
                return set_.begin();
            }

        private:
            Arena arena_;
            melon::F14FastSet<std::string_view> set_;
        };

        template<typename TInput>
        struct InFunctionOuter {
            template<typename TExecCtx>
            struct InFunctionInner {
                POLLUX_DEFINE_FUNCTION_TYPES(TExecCtx);

                MELON_ALWAYS_INLINE void initialize(
                    const std::vector<TypePtr> & /*inputTypes*/,
                    const core::QueryConfig & /*config*/,
                    const arg_type<TInput> * /*searchTerm*/,
                    const arg_type<pollux::Array<TInput> > *searchElements) {
                    if (searchElements == nullptr) {
                        return;
                    }

                    elements_.reserve(searchElements->size());

                    for (const auto &entry: *searchElements) {
                        if (!entry.has_value()) {
                            hasNull_ = true;
                            continue;
                        }
                        elements_.emplace(entry.value());
                    }
                }

                MELON_ALWAYS_INLINE bool callNullable(
                    bool &result,
                    const arg_type<TInput> *searchTerm,
                    const arg_type<pollux::Array<TInput> > * /*array*/) {
                    if (searchTerm == nullptr) {
                        return false;
                    }

                    result = elements_.contains(*searchTerm);
                    if (hasNull_ && !result) {
                        return false;
                    }
                    return true;
                }

            private:
                Set<arg_type<TInput> > elements_;
                bool hasNull_{false};
            };

            template<typename T>
            using Inner = typename InFunctionOuter<TInput>::template InFunctionInner<T>;
        };

        template<typename T>
        void registerInFn(const std::string &prefix) {
            register_function<InFunctionOuter<T>::template Inner, bool, T, Array<T> >(
                {prefix + "in"});
        }
    } // namespace

    void registerIn(const std::string &prefix) {
        registerInFn<int8_t>(prefix);
        registerInFn<int16_t>(prefix);
        registerInFn<int32_t>(prefix);
        registerInFn<int64_t>(prefix);
        registerInFn<float>(prefix);
        registerInFn<double>(prefix);
        registerInFn<bool>(prefix);
        registerInFn<Varchar>(prefix);
        registerInFn<Timestamp>(prefix);
        registerInFn<Date>(prefix);
        registerInFn<ShortDecimal<P1, S1> >(prefix);
        registerInFn<LongDecimal<P1, S1> >(prefix);
    }
} // namespace kumo::pollux::functions::sparksql
